import logging
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, List
from langchain_openai import ChatOpenAI
from langchain.schema import BaseMessage, HumanMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from config import Config
from datetime import datetime

logger = logging.getLogger(__name__)

# Langfuse 可用性检查 - 延迟导入
LANGFUSE_AVAILABLE = None

def _check_langfuse_available():
    """检查Langfuse是否可用"""
    global LANGFUSE_AVAILABLE
    if LANGFUSE_AVAILABLE is None:
        try:
            import langfuse
            LANGFUSE_AVAILABLE = True
        except ImportError:
            LANGFUSE_AVAILABLE = False
    return LANGFUSE_AVAILABLE

class BaseLLMModel(ABC):
    """LLM模型基类"""
    
    def __init__(self, model_name: str = None, temperature: float = 0.7):
        self.model_name = model_name or Config.LLM_MODEL_NAME
        self.temperature = temperature
        
        # 创建LLM，使用LangChain回调系统集成Langfuse
        from monitoring.langfuse_manager import langfuse_manager
        
        # 创建Langfuse回调处理器
        langfuse_handler = None
        if langfuse_manager.is_enabled():
            langfuse_handler = langfuse_manager.create_handler(
                trace_name=f"llm_call_{self.model_name}",
                user_id=None,  # 将在调用时设置
                session_id=None  # 将在调用时设置
            )
        
        # 创建LLM，添加Langfuse回调
        callbacks = []
        if langfuse_handler:
            callbacks.append(langfuse_handler)
        
        self.llm = ChatOpenAI(
            model=self.model_name,
            openai_api_key=Config.OPENAI_API_KEY,
            openai_api_base=Config.OPENAI_BASE_URL,
            temperature=self.temperature,
            callbacks=callbacks
        )
        self._setup_chain()
    
    def _setup_chain(self):
        """设置LangChain链，子类可以重写"""
        pass
    
    def _create_messages(self, system_prompt: str, user_prompt: str) -> List[BaseMessage]:
        """创建消息列表"""
        return [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_prompt)
        ]
    
    def _invoke_with_langfuse(self, messages: List[BaseMessage], **kwargs):
        """使用Langfuse追踪调用LLM - 回调模式"""
        # 提取Langfuse相关参数，避免传递给LLM
        user_id = kwargs.pop("user_id", None)
        session_id = kwargs.pop("session_id", None)
        
        logger.info(f"🔧 BaseLLMModel._invoke_with_langfuse 开始")
        logger.info(f"👤 接收到的追踪参数 - user_id: {user_id}, session_id: {session_id}")
        logger.info(f"📝 消息数量: {len(messages)}")
        
        # 如果LLM已经有Langfuse回调，更新回调中的user_id和session_id
        if hasattr(self.llm, 'callbacks') and self.llm.callbacks:
            logger.info(f"🔄 找到 {len(self.llm.callbacks)} 个回调")
            for i, callback in enumerate(self.llm.callbacks):
                logger.info(f"  回调 {i+1}: {type(callback).__name__}")
                
                if hasattr(callback, 'user_id') and user_id:
                    old_user_id = callback.user_id
                    callback.user_id = user_id
                    logger.info(f"    ✅ 更新user_id: {old_user_id} -> {user_id}")
                if hasattr(callback, 'session_id') and session_id:
                    old_session_id = callback.session_id
                    callback.session_id = session_id
                    logger.info(f"    ✅ 更新session_id: {old_session_id} -> {session_id}")
                    
                # 如果回调有setHandler方法，使用它来设置参数
                if hasattr(callback, 'setHandler'):
                    logger.info(f"    🔧 调用setHandler方法")
                    callback.setHandler(user_id=user_id, session_id=session_id)
        else:
            logger.warning(f"⚠️ 没有找到Langfuse回调")
        
        # 直接调用LLM，Langfuse回调会自动处理追踪
        logger.info(f"🚀 调用LLM.invoke")
        return self.llm.invoke(messages, **kwargs)
    
    def invoke_chain_with_langfuse(self, chain, input_data, **kwargs):
        """使用Langfuse追踪调用LangChain链 - 手动追踪模式"""
        # 提取Langfuse相关参数，避免传递给chain
        user_id = kwargs.pop("user_id", None)
        session_id = kwargs.pop("session_id", None)
        
        # 添加调试日志
        logger.info(f"🔍 invoke_chain_with_langfuse 开始")
        logger.info(f"📋 接收到的参数 - user_id: {user_id}, session_id: {session_id}")
        logger.info(f"📝 input_data: {str(input_data)[:200]}...")
        
        # 使用Langfuse Handler进行手动追踪
        from monitoring.langfuse_manager import langfuse_manager
        import time
        
        if langfuse_manager.is_enabled():
            # 创建Handler
            logger.info(f"🔧 创建Langfuse Handler...")
            handler = langfuse_manager.create_handler(
                trace_name="chain_invoke",
                user_id=user_id,
                session_id=session_id
            )
            
            if handler:
                logger.info(f"✅ Handler创建成功，user_id: {handler.user_id}, session_id: {handler.session_id}")
            else:
                logger.warning(f"⚠️ Handler创建失败")
            
            if handler:
                # 开始追踪
                handler.start_trace()
                
                # 记录开始时间
                start_time = time.time()
                
                # 开始span
                handler.start_span(
                    name=f"chain_{type(chain).__name__}",
                    input_data={
                        "chain_type": type(chain).__name__,
                        "input_data": str(input_data)[:500]  # 限制长度
                    }
                )
                
                try:
                    # 确保chain使用我们的Langfuse handler
                    # 检查chain是否有callbacks属性，如果有则添加我们的handler
                    if hasattr(chain, 'callbacks') and chain.callbacks is None:
                        chain.callbacks = [handler]
                    elif hasattr(chain, 'callbacks') and isinstance(chain.callbacks, list):
                        # 如果已经有callbacks，确保我们的handler在其中
                        if handler not in chain.callbacks:
                            chain.callbacks.append(handler)
                    
                    logger.info(f"🔧 使用Langfuse handler调用chain")
                    logger.info(f"📋 chain callbacks: {chain.callbacks if hasattr(chain, 'callbacks') else 'None'}")
                    
                    # 在调用chain之前，确保LLM的回调也包含正确的user_id和session_id
                    # 检查chain中的LLM组件
                    logger.info(f"🔍 检查chain结构...")
                    logger.info(f"📋 chain类型: {type(chain).__name__}")
                    logger.info(f"📋 chain属性: {dir(chain)}")
                    
                    if hasattr(chain, 'llm'):
                        logger.info(f"✅ 找到chain.llm: {type(chain.llm).__name__}")
                        if hasattr(chain.llm, 'callbacks'):
                            logger.info(f"✅ 找到chain.llm.callbacks: {chain.llm.callbacks}")
                            logger.info(f"📋 LLM回调数量: {len(chain.llm.callbacks)}")
                            for i, callback in enumerate(chain.llm.callbacks):
                                logger.info(f"  回调 {i+1}: {type(callback).__name__}")
                                logger.info(f"    当前user_id: {getattr(callback, 'user_id', 'N/A')}")
                                logger.info(f"    当前session_id: {getattr(callback, 'session_id', 'N/A')}")
                                
                                if hasattr(callback, 'setHandler'):
                                    callback.setHandler(user_id=user_id, session_id=session_id)
                                    logger.info(f"    ✅ 通过setHandler更新LLM回调: {type(callback).__name__}")
                                elif hasattr(callback, 'user_id') and hasattr(callback, 'session_id'):
                                    # 直接设置属性
                                    old_user_id = callback.user_id
                                    old_session_id = callback.session_id
                                    callback.user_id = user_id
                                    callback.session_id = session_id
                                    logger.info(f"    ✅ 直接更新LLM回调属性: {type(callback).__name__}")
                                    logger.info(f"      user_id: {old_user_id} -> {callback.user_id}")
                                    logger.info(f"      session_id: {old_session_id} -> {callback.session_id}")
                        else:
                            logger.warning(f"⚠️ chain.llm没有callbacks属性")
                    else:
                        logger.warning(f"⚠️ chain没有llm属性")
                        
                        # 尝试其他可能的位置
                        if hasattr(chain, 'steps'):
                            logger.info(f"🔍 检查chain.steps...")
                            for i, step in enumerate(chain.steps):
                                logger.info(f"  步骤 {i+1}: {type(step).__name__}")
                                if hasattr(step, 'llm') and hasattr(step.llm, 'callbacks'):
                                    logger.info(f"    ✅ 在步骤中找到LLM回调")
                                    for callback in step.llm.callbacks:
                                        if hasattr(callback, 'setHandler'):
                                            callback.setHandler(user_id=user_id, session_id=session_id)
                                            logger.info(f"      ✅ 更新步骤LLM回调: {type(callback).__name__}")
                    
                    # 直接使用handler调用chain，但不传递config，让chain使用自己的LLM
                    # 在调用之前，确保BaseLLMModel中的LLM回调也被更新
                    logger.info(f"🔧 更新BaseLLMModel中的LLM回调")
                    if hasattr(self, 'llm') and hasattr(self.llm, 'callbacks'):
                        logger.info(f"📋 BaseLLMModel.llm.callbacks数量: {len(self.llm.callbacks)}")
                        for i, callback in enumerate(self.llm.callbacks):
                            logger.info(f"  BaseLLMModel回调 {i+1}: {type(callback).__name__}")
                            if hasattr(callback, 'setHandler'):
                                callback.setHandler(user_id=user_id, session_id=session_id)
                                logger.info(f"    ✅ 更新BaseLLMModel回调: {type(callback).__name__}")
                    
                    result = chain.invoke(input_data, **kwargs)
                    
                    # 计算执行时间
                    end_time = time.time()
                    latency = end_time - start_time
                    
                    # 结束span
                    handler.end_span(output={
                        "result": str(result)[:500],  # 限制长度
                        "latency": latency
                    })
                    
                    # 结束追踪
                    handler.end_trace(output={
                        "status": "success",
                        "latency": latency
                    })
                    
                    return result
                except Exception as e:
                    # 计算执行时间
                    end_time = time.time()
                    latency = end_time - start_time
                    
                    # 结束span（错误）
                    handler.end_span(output={
                        "error": str(e),
                        "latency": latency
                    })
                    
                    # 结束追踪
                    handler.end_trace(output={
                        "status": "error", 
                        "error": str(e),
                        "latency": latency
                    })
                    raise
            else:
                # Handler创建失败，直接调用
                return chain.invoke(input_data, **kwargs)
        else:
            # 不使用Langfuse追踪
            return chain.invoke(input_data, **kwargs)
    
    @abstractmethod
    def generate_response(self, **kwargs) -> str:
        """生成响应的抽象方法"""
        pass
    
    def _add_metadata(self, result: str, **kwargs) -> Dict[str, Any]:
        """添加元数据到结果中"""
        return {
            "result": result,
            "model_used": self.model_name,
            "timestamp": datetime.now().isoformat(),
            **kwargs
        }
    
    def _handle_error(self, error: Exception, operation: str = "操作") -> str:
        """统一错误处理"""
        return f"{operation}时发生错误: {str(error)}" 